package cymbol;

import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.misc.MultiMap;
import org.antlr.v4.runtime.misc.OrderedHashSet;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.ParseTreeWalker;
import org.stringtemplate.v4.ST;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Set;

public class CallGraph {

    // simple graph model of the output.
    // eg: f ->[g, h]
    static class Graph {
        Set<String> nodes = new OrderedHashSet<>();
        MultiMap<String, String> edges = new MultiMap<>();

        public void edge(String source, String target) {
            edges.map(source, target);
        }

        public String toString() {
            return "edges: " + edges.toString() + ", functions: " + nodes;
        }

        public String toDOT() {
            StringBuilder buf = new StringBuilder();
            buf.append("digraph G{\n");
            buf.append(" ranksep=.25;\n");
            buf.append(" edge [arrowsize=.5]\n");
            buf.append(" node [shape=circle, fontname=\"ArialNarrow\",\n");
            buf.append("      fontsize=12, fixedsize=true, height=.45];\n");
            buf.append(" ");
            // print all nodes first
            for (String node : nodes) {
                buf.append(node).append("; ");
            }
            buf.append("\n");
            for (String src : edges.keySet()) {
                for (String trg : edges.get(src)) {
                    buf.append(" ").append(src).append(" -> ").append(trg).append(";\n");
                }
            }
            buf.append("}\n");
            return buf.toString();
        }

       /* fill stringTemplate:
        digraph G {
            rankdir=LR;
            <edgePairs: {edge | <edge.a> -> <edge.b>; }; separator="\n">
            <childless: {f | <f>; }; separator="\n">
        }
        */

        public ST toST() {
            ST st = new ST(
                    "digraph G{\n" +
                            " rankesep=.25;\n" +
                            " edge [arrowsize=.5]\n" +
                            " node [shape=circle, fontname=\"ArialNarrow\",\n" +
                            "       fontsize=12, fixedsize=true, height=.45];\n" +
                            " <funcs: {f | <f>; }>\n" +
                            " <edgePairs: {edge | <edge.a> -> <edge.b>;}; separator=\"\\n\">\n" +
                            "}\n"
            );
            st.add("edgePairs", edges.getPairs());
            st.add("funcs", nodes);
            return st;
        }
    }

    static class FunctionListener extends CymbolBaseListener {
        Graph graph = new Graph();
        String currentFunctionName = null;

        public void enterFunctionDecl(CymbolParser.FunctionDeclContext context) {
            currentFunctionName = context.ID().getText();
            graph.nodes.add(currentFunctionName);
        }

        public void exitCall(CymbolParser.CallContext context) {
            String funcName = context.ID().getText();
            graph.edge(currentFunctionName, funcName);
        }
    }

    public static void main(String[] args) throws IOException {
        String inputFile = null;
        if (args.length > 0) {
            inputFile = args[0];
        }
        InputStream is = System.in;
        if (inputFile != null) {
            is = new FileInputStream(inputFile);
        }
        CharStream input = CharStreams.fromStream(is);
        CymbolLexer lexer = new CymbolLexer(input);
        CommonTokenStream tokenStream = new CommonTokenStream(lexer);
        CymbolParser parser = new CymbolParser(tokenStream);
        parser.setBuildParseTree(true);
        ParseTree tree = parser.file();

        ParseTreeWalker walker = new ParseTreeWalker();
        FunctionListener collector = new FunctionListener();
        walker.walk(collector, tree);
        System.out.println(collector.graph.toString());
        System.out.println(collector.graph.toDOT());
        // using stringTemplate to generate output
        System.out.println("using string template:\n");
        System.out.println(collector.graph.toST().render());
    }
}
